// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <gtest/gtest.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/core/plan_node.h>

using namespace ::kumo::pollux;
using namespace ::kumo::pollux::core;

TEST(TestPlanNode, findFirstNode) {
  auto rowType = ROW({"name1"}, {BIGINT()});

  std::shared_ptr<connector::ConnectorTableHandle> tableHandle;
    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
      assignments;

  std::shared_ptr<PlanNode> tableScan3 =
      std::make_shared<TableScanNode>("3", rowType, tableHandle, assignments);
  std::shared_ptr<PlanNode> tableScan2 =
      std::make_shared<TableScanNode>("2", rowType, tableHandle, assignments);

  std::vector<FieldAccessTypedExprPtr> sortingKeys;
  std::vector<SortOrder> sortingOrders;
  std::shared_ptr<PlanNode> localMerge1 = std::make_shared<LocalMergeNode>(
      "1",
      sortingKeys,
      sortingOrders,
      std::vector<PlanNodePtr>{tableScan2, tableScan3});

  std::vector<std::string> names;
  std::vector<TypedExprPtr> projections;
  std::shared_ptr<PlanNode> project0 =
      std::make_shared<ProjectNode>("0", names, projections, localMerge1);

  EXPECT_EQ(
      tableScan3.get(),
      PlanNode::findFirstNode(project0.get(), [](const PlanNode* node) {
        return node->id() == "3";
      }));

  EXPECT_EQ(
      project0.get(),
      PlanNode::findFirstNode(project0.get(), [](const PlanNode* node) {
        return node->name() == "Project";
      }));

  EXPECT_EQ(
      nullptr,
      PlanNode::findFirstNode(project0.get(), [](const PlanNode* node) {
        return node->name() == "Unknown";
      }));
}

TEST(TestPlanNode, sortOrder) {
  struct {
    SortOrder order1;
    SortOrder order2;
    int expectedEqual;

    std::string debugString() const {
      return fmt::format(
          "order1 {} order2 {} expectedEqual {}",
          order1.toString(),
          order2.toString(),
          expectedEqual);
    }
  } testSettings[] = {
      {{true, true}, {true, true}, true},
      {{true, false}, {true, false}, true},
      {{false, true}, {false, true}, true},
      {{false, false}, {false, false}, true},
      {{true, true}, {true, false}, false},
      {{true, true}, {false, false}, false},
      {{true, true}, {false, true}, false},
      {{true, false}, {false, false}, false},
      {{true, false}, {false, true}, false},
      {{false, true}, {false, false}, false}};
  for (const auto& testData : testSettings) {
    SCOPED_TRACE(testData.debugString());
    if (testData.expectedEqual) {
      ASSERT_EQ(testData.order1, testData.order2);
    } else {
      ASSERT_NE(testData.order1, testData.order2);
    }
  }
}

TEST(TestPlanNode, duplicateSortKeys) {
  auto sortingKeys = std::vector<FieldAccessTypedExprPtr>{
      std::make_shared<core::FieldAccessTypedExpr>(BIGINT(), "c0"),
      std::make_shared<core::FieldAccessTypedExpr>(BIGINT(), "c1"),
      std::make_shared<core::FieldAccessTypedExpr>(BIGINT(), "c0"),
  };
  auto sortingOrders =
      std::vector<SortOrder>{{true, true}, {false, false}, {true, true}};
  POLLUX_ASSERT_USER_THROW(
      std::make_shared<OrderByNode>(
          "orderBy", sortingKeys, sortingOrders, false, nullptr),
      "Duplicate sorting keys are not allowed: c0");
}
